#include <gtest/gtest.h>
#include <random>
#include <vector>
#include <cmath>
#include <algorithm>

#include "saxpy.h"
#include "gemv.h"
#include "pack.h"

class AccuracyTest : public ::testing::Test {
protected:
    void SetUp() override {
        gen.seed(20250526);
    }

    std::mt19937 gen;
    std::uniform_real_distribution<float> float_dist{-1000.0f, 1000.0f};
    std::uniform_int_distribution<uint8_t> uint8_dist{0, 15};
    std::uniform_int_distribution<size_t> size_dist{100, 10000};
};

bool compare_arrays(const float* a, const float* b, size_t n, float tolerance = 1e-6f) {
    // Calculate L2 norm of the difference
    double diff_norm_sq = 0.0;
    double b_norm_sq = 0.0;
    
    for (size_t i = 0; i < n; i++) {
        double diff = static_cast<double>(a[i]) - static_cast<double>(b[i]);
        diff_norm_sq += diff * diff;
        b_norm_sq += static_cast<double>(b[i]) * static_cast<double>(b[i]);
    }
    
    // Calculate relative L2 error: ||a - b||_2 / ||b||_2
    double diff_norm = std::sqrt(diff_norm_sq);
    double b_norm = std::sqrt(b_norm_sq);
    
    // Avoid division by zero
    if (b_norm < 1e-15) {
        // If b is essentially zero, check if a is also essentially zero
        bool result = diff_norm < tolerance;
        if (!result) {
            std::cerr << "Arrays are not equal: ||a - b||_2 = " << diff_norm 
                      << ", ||b||_2 ≈ 0, tolerance = " << tolerance << std::endl;
        }
        return result;
    }
    
    double relative_error = diff_norm / b_norm;
    
    if (relative_error > tolerance) {
        std::cerr << "Arrays are not equal: relative L2 error = " << relative_error 
                  << " (||a - b||_2 = " << diff_norm << ", ||b||_2 = " << b_norm 
                  << "), tolerance = " << tolerance << std::endl;
        return false;
    }
    
    return true;
}

bool compare_arrays(const uint8_t* a, const uint8_t* b, size_t n) {
    for (size_t i = 0; i < n; i++) {
        if (a[i] != b[i]) {
            std::cerr << "Arrays are not equal at index " << static_cast<int>(i) << ", expected " << static_cast<int>(a[i]) << " but got " << static_cast<int>(b[i]) << std::endl;
            return false;
        }
    }
    return true;
}

// SAXPY Tests
TEST_F(AccuracyTest, SaxpyRandomSmallInputs) {
    const int num_tests = 1000;
    
    for (int test = 0; test < num_tests; test++) {
        size_t n = size_dist(gen) % 1000 + 10;
        float a = float_dist(gen);
        
        std::vector<float> x(n), y_scalar(n), y_sve(n);
        
        // Generate random inputs
        for (size_t i = 0; i < n; i++) {
            x[i] = float_dist(gen);
            y_scalar[i] = y_sve[i] = float_dist(gen);
        }
        
        // Run both implementations
        saxpy(x.data(), y_scalar.data(), a, n);
        saxpy_sve(x.data(), y_sve.data(), a, n);
        
        // Compare results
        ASSERT_TRUE(compare_arrays(y_scalar.data(), y_sve.data(), n));
    }
}

TEST_F(AccuracyTest, SaxpyRandomLargeInputs) {
    const int num_tests = 100;
    
    for (int test = 0; test < num_tests; test++) {
        size_t n = size_dist(gen) + 10000; // Large arrays
        float a = float_dist(gen);
        
        std::vector<float> x(n), y_scalar(n), y_sve(n);
        
        // Generate random inputs
        for (size_t i = 0; i < n; i++) {
            x[i] = float_dist(gen);
            y_scalar[i] = y_sve[i] = float_dist(gen);
        }
        
        // Run both implementations
        saxpy(x.data(), y_scalar.data(), a, n);
        saxpy_sve(x.data(), y_sve.data(), a, n);
        
        // Compare results
        ASSERT_TRUE(compare_arrays(y_scalar.data(), y_sve.data(), n));
    }
}

TEST_F(AccuracyTest, SaxpyEdgeCases) {
    // Test with zero scalar
    size_t n = 100;
    std::vector<float> x(n), y_scalar(n), y_sve(n);
    
    for (size_t i = 0; i < n; i++) {
        x[i] = float_dist(gen);
        y_scalar[i] = y_sve[i] = float_dist(gen);
    }
    
    saxpy(x.data(), y_scalar.data(), 0.0f, n);
    saxpy_sve(x.data(), y_sve.data(), 0.0f, n);
    
    ASSERT_TRUE(compare_arrays(y_scalar.data(), y_sve.data(), n));
    
    // Test with very small arrays
    for (size_t small_n = 1; small_n <= 10; small_n++) {
        std::vector<float> x_small(small_n), y_scalar_small(small_n), y_sve_small(small_n);
        float a = float_dist(gen);
        
        for (size_t i = 0; i < small_n; i++) {
            x_small[i] = float_dist(gen);
            y_scalar_small[i] = y_sve_small[i] = float_dist(gen);
        }
        
        saxpy(x_small.data(), y_scalar_small.data(), a, small_n);
        saxpy_sve(x_small.data(), y_sve_small.data(), a, small_n);
        
        ASSERT_TRUE(compare_arrays(y_scalar_small.data(), y_sve_small.data(), small_n));
    }
}

// GEMV Tests
TEST_F(AccuracyTest, GemvRandomInputs) {
    const int num_tests = 500;
    
    for (int test = 0; test < num_tests; test++) {
        size_t m = (size_dist(gen) % 500) + 10;
        size_t n = (size_dist(gen) % 500) + 10;
        
        std::vector<float> A(m * n), x(n), y_scalar(m), y_sve(m);
        
        // Generate random inputs
        for (size_t i = 0; i < m * n; i++) {
            A[i] = float_dist(gen);
        }
        for (size_t i = 0; i < n; i++) {
            x[i] = float_dist(gen);
        }

        // Run both implementations
        gemv(A.data(), x.data(), y_scalar.data(), m, n);
        gemv_sve(A.data(), x.data(), y_sve.data(), m, n);
        
        // Compare results
        ASSERT_TRUE(compare_arrays(y_scalar.data(), y_sve.data(), m));
    }
}

TEST_F(AccuracyTest, GemvSquareMatrices) {
    const int num_tests = 200;
    
    for (int test = 0; test < num_tests; test++) {
        size_t n = (size_dist(gen) % 200) + 10; // Square matrices
        
        std::vector<float> A(n * n), x(n), y_scalar(n), y_sve(n);
        
        // Generate random inputs
        for (size_t i = 0; i < n * n; i++) {
            A[i] = float_dist(gen);
        }
        for (size_t i = 0; i < n; i++) {
            x[i] = float_dist(gen);
        }
        
        // Run both implementations
        gemv(A.data(), x.data(), y_scalar.data(), n, n);
        gemv_sve(A.data(), x.data(), y_sve.data(), n, n);
        
        // Compare results
        ASSERT_TRUE(compare_arrays(y_scalar.data(), y_sve.data(), n));
    }
}

TEST_F(AccuracyTest, GemvEdgeCases) {
    // Test with small matrices
    for (size_t m = 1; m <= 8; m++) {
        for (size_t n = 1; n <= 8; n++) {
            std::vector<float> A(m * n), x(n), y_scalar(m), y_sve(m);
            
            for (size_t i = 0; i < m * n; i++) {
                A[i] = float_dist(gen);
            }
            for (size_t i = 0; i < n; i++) {
                x[i] = float_dist(gen);
            }
            
            gemv(A.data(), x.data(), y_scalar.data(), m, n);
            gemv_sve(A.data(), x.data(), y_sve.data(), m, n);
            
            ASSERT_TRUE(compare_arrays(y_scalar.data(), y_sve.data(), m));
        }
    }
}

// Pack/Unpack Tests
TEST_F(AccuracyTest, PackUnpackRandomInputs) {
    const int num_tests = 1000;
    
    for (int test = 0; test < num_tests; test++) {
        size_t n = ((size_dist(gen) % 1000) + 10) & ~1;

        std::vector<uint8_t> src(n), dst_scalar(n / 2), dst_sve(n / 2);
        std::vector<uint8_t> unpacked_scalar(n), unpacked_sve(n);
        
        // Generate random 4-bit inputs
        for (size_t i = 0; i < n; i++) {
            src[i] = uint8_dist(gen);
        }
        
        // Test pack functions
        pack(dst_scalar.data(), src.data(), n);
        pack(dst_sve.data(), src.data(), n);

        // ASSERT_TRUE(compare_arrays(dst_scalar.data(), dst_sve.data(), n / 2))
        //     << "Pack mismatch at test " << test << " with n=" << n;

        // Test unpack functions
        unpack(unpacked_scalar.data(), dst_scalar.data(), n);
        unpack_sve(unpacked_sve.data(), dst_scalar.data(), n);

        ASSERT_TRUE(compare_arrays(unpacked_scalar.data(), unpacked_sve.data(), n))
            << "Unpack mismatch at test " << test << " with n=" << n;

        // Test round-trip: original -> pack -> unpack should equal original
        ASSERT_TRUE(compare_arrays(src.data(), unpacked_scalar.data(), n))
            << "Pack/Unpack round-trip failed at test " << test;
    }
}

TEST_F(AccuracyTest, PackUnpackEdgeCases) {
    // Test with minimum size
    size_t n = 2;
    std::vector<uint8_t> src = {0xA, 0xB};
    std::vector<uint8_t> dst_sve(2);
    std::vector<uint8_t> unpacked_sve(2);

    pack(dst_sve.data(), src.data(), n);
    
    // Expected: (0xA << 4) | 0xB = 0xAB
    ASSERT_EQ(dst_sve[0], 0xAB);

    unpack_sve(unpacked_sve.data(), dst_sve.data(), n);
    
    ASSERT_EQ(unpacked_sve[0], 0xA);
    ASSERT_EQ(unpacked_sve[1], 0xB);
    
    // Test with various patterns
    std::vector<uint8_t> patterns = {0x0, 0xF, 0x5, 0xA};
    for (size_t i = 0; i < patterns.size(); i++) {
        for (size_t j = 0; j < patterns.size(); j++) {
            std::vector<uint8_t> test_src = {patterns[i], patterns[j]};
            std::vector<uint8_t> test_dst(1), test_unpacked(2);
            
            pack(test_dst.data(), test_src.data(), 2);
            unpack(test_unpacked.data(), test_dst.data(), 2);
            
            ASSERT_EQ(test_unpacked[0], patterns[i]);
            ASSERT_EQ(test_unpacked[1], patterns[j]);
        }
    }
}
